"""
python fujian2_LinearRegression_fillBlank.py
"""

import os
import json
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from datetime import datetime, timedelta
from sklearn.linear_model import LinearRegression

def LinearRegression_fillBlank_single_category(input_path, output_path):
    # 1. 加载数据
    with open(input_path, 'r', encoding='utf-8') as f:
        data = json.load(f)

    # 将数据转为DataFrame格式
    df = pd.DataFrame(data)
    df['date'] = pd.to_datetime(df['date'])
    df.set_index('date', inplace=True)

    # 2. 确定需要填补的日期范围
    fill_start_date = pd.to_datetime("2022-10-01")
    fill_end_date = pd.to_datetime("2023-03-31")
    fill_dates = pd.date_range(start=fill_start_date, end=fill_end_date)

    # 3. 构建线性回归模型进行填补
    # 选择已知的日期和销量
    known_data = df[(df.index >= '2022-07-01') & (df.index <= '2022-09-30')]
    known_data = pd.concat([known_data, df[(df.index >= '2023-04-01') & (df.index <= '2023-06-30')]])

    # 将日期转换为数值形式
    known_data['date_ordinal'] = known_data.index.map(datetime.toordinal)
    
    # 准备训练数据
    X = known_data['date_ordinal'].values.reshape(-1, 1)
    y = known_data['sales'].values

    # 训练线性回归模型
    model = LinearRegression()
    model.fit(X, y)

    # 4. 预测空白期的销量
    fill_dates_ordinal = fill_dates.map(datetime.toordinal).values.reshape(-1, 1)
    predicted_sales = model.predict(fill_dates_ordinal)

    # 5. 保存填补后的数据
    filled_data = [{'date': date.strftime('%Y/%m/%d'), 'sales': int(sales)} 
                   for date, sales in zip(fill_dates, predicted_sales)]

    # 检查并创建输出目录
    output_dir = os.path.dirname(output_path)
    os.makedirs(output_dir, exist_ok=True)

    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(filled_data, f, ensure_ascii=False, indent=4)

    print(f"填补后的数据已保存至 {output_path}")

    # 6. 绘制图形
    plt.figure(figsize=(14, 7))
    
    # 原始数据
    plt.scatter(df.index, df['sales'], color='blue', label='origin data')

    # 填补的数据
    fill_df = pd.DataFrame(filled_data)
    fill_df['date'] = pd.to_datetime(fill_df['date'])
    plt.plot(fill_df['date'], fill_df['sales'], color='orange', label='filled data')

    plt.title('销量数据及填补')
    plt.xlabel('日期')
    plt.ylabel('销量')
    plt.legend()
    plt.savefig(os.path.splitext(output_path)[0] + '_plot.png')
    plt.close()
    
    print(f"图像已保存至 {os.path.splitext(output_path)[0] + '_plot.png'}")

if __name__ == '__main__':
    input_path = '../fujian/fujian2/groupByCategory/category_category2.json'
    output_path = '../fujian/fujian2/LinearRegression/category_category2_filled.json'

    LinearRegression_fillBlank_single_category(input_path, output_path)
